/*
    Author: Christopher Diggins
    License: MIT Licence 1.0
    Website: http://www.heron-language.com
    
    A typed AST generated from the untyped AST create by the YARD Heron parser.
*/

#ifndef YARD_HERON_AST_HPP
#define YARD_HERON_AST_HPP

#include <vector>
#include <string>

namespace heron_ast
{
    typedef Node YardNode;

    //////////////////////////////////////////////////////////////////////////
    // base classes 

    struct AstException
    {
        const char* msg;
        AstException(const char* _msg)
            : msg(_msg)
        { }
    };
    
    template<typename T>
    struct AstNodeList
    {
        void Add(T* x) {
            nodes.push_back(x);
        };
        int Count() {
            return nodes.size();
        }
        T* At(int i) {
            return nodes[i];
        }
        template<typename F>
        void WithEach(F& f) {
            for (int i=0; i < Count(); ++i) {
                f(*At(i));
            }
        }
    private:
        std::vector<T*> nodes;
    };

    class AstNode
    {
        YardNode* node;
        AstNode* parent;
        std::vector<AstNode*> children;

        // This field is used only for construction 
        // the better alternative may be to have each "Create" function
        // pass a state object to each other.
        // I prefer the simplicity of the approach of hiding construction
        // details in the class itself.
        YardNode* cur; 
 
    public:

        AstNode(AstNode* _parent, YardNode* _node)
            : node(_node), parent(_parent), cur(NULL)
        {
            if (node->HasChildren())
                cur = node->GetFirstChild();
        }

        virtual ~AstNode() 
        {
            for (int i=0; i < GetNumChildren(); ++i) {
                delete(GetChild(i));
                children[i] = NULL;
            }
        }

        template<typename T>
        bool TypeMatches() {
            return (typeid(*this) == typeid(T));
        }

        template<typename T, typename F>
        void WithEachTypedChild(F& f) {
            for (int i=0; i < GetNumChildren(); ++i) {
                if (GetChild(i)->TypeMatches<T>()) {
                    f(GetChild(i)->CastTo<T>());
                }
            }
        }

        template<typename F>
        void WithEachChild(F& f) {
            for (int i=0; i < GetNumChildren(); ++i) {
                f(GetChild(i));
            }
        }

        template<typename T>
        T* CastTo() {
            if (!TypeMatches<T>()) {
                AstError<T>("bad cast");
            }
            return dynamic_cast<T*>(this);
        }

        int GetNumChildren() 
        {
            return children.size();
        }

        AstNode* GetChild(int i) 
        {
            return children[i];
        }

        template<typename T>
        void AstError(const char* x)
        {
            throw AstException(x);
        }

        template<typename T>
        void Create(T*& x) 
        {
            if (cur == NULL) {
                AstError<T>("no children nodes");
                return;
            }
            if (!cur->TypeMatches<typename T::Rule>()) {
                AstError<T>("node rule does not match type");
                return;
            }
            x = new T(this, cur);
            cur = cur->GetSibling();
            children.push_back(x);
        }

        template<typename T>
        bool CreateIf(AstNode*& x) 
        {
            if (cur == NULL) {
                AstError<T>("no children nodes");
                return false;
            }
            if (!cur->TypeMatches<typename T::Rule>()) {
                return false;
            }
            x = new T(this, cur);
            cur = cur->GetSibling();
            children.push_back(x);
            return true;
        }

        template<typename T>
        void CreateOpt(T*& x)
        {
            if (cur == NULL) {
                x = NULL;
                return;
            }
            if (!cur->TypeMatches<typename T::Rule>()) {
                x = NULL;
                return;
            }
            x = new T(this, cur);
            cur = cur->GetSibling();
            children.push_back(x);
        }

        template<typename T>
        void Create(AstNodeList<T>& x) 
        {
            while (cur != NULL && cur->TypeMatches<typename T::Rule>()) {
                T* tmp = new T(this, cur);
                x.Add(tmp);
                children.push_back(tmp);
                cur = cur->GetSibling();
            }
        }

        std::string ToString() {
            return Trim(std::string(node->Begin(), node->End()));
        }
    };
    
    //////////////////////////////////////////////////////////////////////////
    // specializations 

    struct Expr;
    struct TypeExpr;
    struct SimpleExpr;
    struct Statement;
    struct StatementList;
    struct CodeBlock;

    struct Name : AstNode
    {
        typedef heron_grammar::Name Rule;

        Name(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
        }
    };

    struct BinNumber : AstNode
    {
        typedef heron_grammar::BinNumber Rule;

        BinNumber(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
        }
    };

    struct HexNumber : AstNode
    {
        typedef heron_grammar::HexNumber Rule;

        HexNumber(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
        }
    };

    struct DecNumber : AstNode
    {
        typedef heron_grammar::DecNumber Rule;

        DecNumber(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
        }
    };

    struct CharLiteral : AstNode
    {
        typedef heron_grammar::CharLiteral Rule;

        CharLiteral(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
        }
    };

    struct StringLiteral : AstNode
    {
        typedef heron_grammar::StringLiteral Rule;

        StringLiteral(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
        }
    };

    struct Literal : AstNode
    {
        typedef heron_grammar::Literal Rule;

        AstNode* value;

        Literal(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            if (CreateIf<BinNumber>(value)
                || CreateIf<HexNumber>(value)
                || CreateIf<DecNumber>(value)
                || CreateIf<CharLiteral>(value)
                || CreateIf<StringLiteral>(value)
                )
                return;
            throw new AstException("unrecognized literal type");
        }
    };

    struct TypeArgs : AstNode
    {
        typedef heron_grammar::TypeArgs Rule;

        AstNodeList<TypeExpr> args;

        TypeArgs(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(args);
        }
    };

    struct TypeExpr : AstNode
    {
        typedef heron_grammar::TypeExpr Rule;

        Name* name;
        TypeArgs* typeargs;

        TypeExpr(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(name);
            CreateOpt(typeargs);
        }
    };

    struct Arg : AstNode
    {
        typedef heron_grammar::Arg Rule;

        Name* name;
        TypeExpr* type;

        Arg(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(name);
            CreateOpt(type);
        }
    };

    struct ArgList : AstNode
    {
        typedef heron_grammar::ArgList Rule;

        AstNodeList<Arg> args;

        ArgList(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(args);
        }
    };

    struct AnonFxn : AstNode
    {
        typedef heron_grammar::AnonFxn Rule;

        ArgList* expr;
        TypeExpr* type;
        CodeBlock* body;

        AnonFxn(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(expr);
            CreateOpt(type);
            Create(body);
        }
    };

    struct DelExpr : AstNode
    {
        typedef heron_grammar::DelExpr Rule;

        Expr* expr;

        DelExpr(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(expr);
        }
    };

    struct ParanthesizedExpr : AstNode
    {
        typedef heron_grammar::ParanthesizedExpr Rule;

		AstNodeList<Expr> subexprs;        

        ParanthesizedExpr(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(subexprs);
        }
    };

    struct BracketedExpr : AstNode
    {
        typedef heron_grammar::BracketedExpr Rule;

        Expr* expr;

        BracketedExpr(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(expr);
        }
    };

    struct NewExpr : AstNode
    {
        typedef heron_grammar::NewExpr Rule;

        TypeExpr* type;
        ParanthesizedExpr* args;

        NewExpr(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(type);
            Create(args);
        }
    };

    struct SimpleExpr : AstNode
    {
        typedef heron_grammar::SimpleExpr Rule;

        AstNode* expr; 

        SimpleExpr(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            if (CreateIf<NewExpr>(expr)
                || CreateIf<DelExpr>(expr)
                || CreateIf<Name>(expr)
                || CreateIf<Literal>(expr)
                || CreateIf<AnonFxn>(expr)
                || CreateIf<ParanthesizedExpr>(expr)
                || CreateIf<BracketedExpr>(expr)
                )
                return;
            throw new AstException("unrecognized literal type");
        }
    };

    struct Expr : AstNode
    {
        typedef heron_grammar::Expr Rule;

        AstNodeList<SimpleExpr> exprs;

        Expr(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(exprs);
        }
    };

    struct CodeBlock : AstNode
    {
        typedef heron_grammar::CodeBlock Rule;

        AstNodeList<Statement> statements;

        CodeBlock(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(statements);
        }
    };

    struct VarDecl : AstNode
    {
        typedef heron_grammar::VarDecl Rule;

        Name* name;
        TypeExpr* type;
        Expr* init;
    
        VarDecl(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(name);
            CreateOpt(type);
            CreateOpt(init);
        }
    };

    struct StaticDecl : AstNode
    {
        typedef heron_grammar::StaticDecl Rule;

        Name* name;
        TypeExpr* type;
        Expr* init;
    
        StaticDecl(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(name);
            CreateOpt(type);
            CreateOpt(init);
        }
    };

    struct IfStatement : AstNode
    {
        typedef heron_grammar::IfStatement Rule;

        Expr* cond;
        Statement* body;
        Statement* alternative;
    
        IfStatement(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(cond);
            Create(body);
            CreateOpt(alternative);
        }
    };

    struct ForStatement : AstNode
    {
        typedef heron_grammar::ForStatement Rule;

        Name* name;
		TypeExpr* type;
		Expr* init;
        Expr* inv;
        Expr* iter;
        Statement* body;
    
        ForStatement(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
			Create(name);
			CreateOpt(type);
			Create(init);
            Create(inv);
            Create(iter);
            Create(body);
        }
    };

    struct ForEachStatement : AstNode
    {
        typedef heron_grammar::ForEachStatement Rule;

        Name* name;
        TypeExpr* type;
        Expr* coll;
        Statement* body;
    
        ForEachStatement(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(name);
            CreateOpt(type);
            Create(coll);
            Create(body);
        }        
    };

    struct ExprStatement : AstNode
    {
        typedef heron_grammar::ExprStatement Rule;

        Expr* expr;

        ExprStatement(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(expr);
        }
    };

    struct ReturnStatement : AstNode
    {
        typedef heron_grammar::ReturnStatement Rule;

        Expr* expr;

        ReturnStatement(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(expr);
        }
    };

    struct WhileStatement : AstNode
    {
        typedef heron_grammar::WhileStatement Rule;

        Expr* cond;
        Statement* body;

        WhileStatement(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        { 
            Create(cond);
            Create(body);
        };
    };

    struct ImportStatement : AstNode
    {
        typedef heron_grammar::ImportStatement Rule;

        Name* name;

        ImportStatement(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(name);
        }
    };

    struct Function : AstNode 
    {
        typedef heron_grammar::Function Rule;

        Name* name;
        ArgList* args;
        TypeExpr* type;
        CodeBlock* body;

        Function(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(name);
            Create(args);
            CreateOpt(type);
            Create(body);
        }
    };

    struct Statement : AstNode
    {
        typedef heron_grammar::Statement Rule;

        AstNode* statement;

        Statement(AstNode* parent, YardNode* node)
            : AstNode(parent, node)
        {
            if (CreateIf<CodeBlock>(statement)
                || CreateIf<VarDecl>(statement)
                || CreateIf<IfStatement>(statement)
                || CreateIf<ForStatement>(statement)
                || CreateIf<ForEachStatement>(statement)
                || CreateIf<WhileStatement>(statement)
                || CreateIf<ReturnStatement>(statement)
                || CreateIf<ExprStatement>(statement)
				|| CreateIf<Function>(statement)
				|| CreateIf<ImportStatement>(statement)
				|| CreateIf<StaticDecl>(statement)
                )
                return;
            throw AstException("unrecognized statement type");
        }
    };

	struct Inherits : AstNode
	{
		typedef heron_grammar::Inherits Rule;

		AstNodeList<TypeExpr> types;

		Inherits(AstNode* parent, YardNode* node) 
			: AstNode(parent, node)
		{
			Create(types);
		}		  
	};

	struct Implements : AstNode
	{
		typedef heron_grammar::Implements Rule;

		AstNodeList<TypeExpr> types;

		Implements(AstNode* parent, YardNode* node) 
			: AstNode(parent, node)
		{
			Create(types);
		}		  
	};

    struct Field : AstNode
    {
        typedef heron_grammar::Field Rule;

        Name* name;
        TypeExpr* type;
    
        Field(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(name);
            CreateOpt(type);
        }
    };

	struct Fields : AstNode
	{
		typedef heron_grammar::Fields Rule;

		AstNodeList<Field> fields;

		Fields(AstNode* parent, YardNode* node) 
			: AstNode(parent, node)
		{
			Create(fields);
		}		  
	};

    struct Method : AstNode 
    {
        typedef heron_grammar::Method Rule;

        Name* name;
        ArgList* args;
        TypeExpr* type;
		CodeBlock* body;

        Method(AstNode* parent, YardNode* node) 
            : AstNode(parent, node)
        {
            Create(name);
            Create(args);
            CreateOpt(type);
			CreateOpt(body);
        }
    };
	
	struct Methods : AstNode
	{
		typedef heron_grammar::Methods Rule;

		AstNodeList<Method> methods;

		Methods(AstNode* parent, YardNode* node) 
			: AstNode(parent, node)
		{
			Create(methods);
		}		  
	};

	struct Class : AstNode
	{
        typedef heron_grammar::Class Rule;

		Name* name;
		Inherits* inherits;
		Implements* implements;
		Fields* fields;
		Methods* methods;

		Class(AstNode* parent, YardNode* node) :
			AstNode(parent, node)
		{
			Create(name);
			CreateOpt(inherits);
			CreateOpt(implements);
			CreateOpt(methods);
			CreateOpt(fields);
		}
	};

	struct Module : AstNode
	{
        typedef heron_grammar::Module Rule;

		Name* name;
		AstNodeList<Class> classes;

		Module(AstNode* parent, YardNode* node) :
			AstNode(parent, node)
		{
			Create(name);
			Create(classes);
		}
	};

	struct Program : AstNode
	{
        typedef heron_grammar::Program Rule;

		Name* name;
		CodeBlock* block;

		Program(AstNode* parent, YardNode* node) :
			AstNode(parent, node)
		{
			Create(name);
			Create(block);
		}
	};

	struct StatementList : AstNode
	{
        typedef heron_grammar::StatementList Rule;

		AstNodeList<Statement> statements;

		StatementList(AstNode* parent, YardNode* node) :
			AstNode(parent, node)
		{
			Create(statements);
		}
	};

    struct HeronFile : AstNode 
    {
        typedef heron_grammar::HeronFile Rule;

		AstNode* program_or_module;

        HeronFile(AstNode* parent, YardNode* node) : 
            AstNode(parent, node)
        {
            if (CreateIf<Program>(program_or_module)
				|| CreateIf<Module>(program_or_module)
				|| CreateIf<StatementList>(program_or_module))
				return;

            throw AstException("ill-formed script file");
        }
    };

	struct HeronAst 
	{
		char* data;
		int datasize; 
		HeronFile* root;
		yard::SimpleTextParser* parser;

		// TODO: should probsably be a stream.
		std::string sErrMsg;

		HeronAst()
			: data(NULL), datasize(0), root(NULL), parser(NULL), sErrMsg("no content loaded\n")
		{
		}

		void Clear() {
			if (root != NULL) delete root;
			root = NULL;
			if (parser != NULL) delete parser;
			parser = NULL;

			// The data is malloc'd
			if (data != NULL) free(data);
			data = NULL;
		}

		~HeronAst() {
			Clear();
		}

		bool LoadFile(const char* sInFile) {
			Clear();
			sErrMsg = "";
			datasize = FileSize(sInFile);			
			if (datasize == 0) {
				sErrMsg += "could not open file: ";
				sErrMsg += sInFile;
				sErrMsg += "\n";
				return false;
			}
			//printf("File size = %d\n", datasize);
			data = ReadFile(sInFile, datasize);
			if (data == NULL) {
				sErrMsg += "could not load file: ";
				sErrMsg += sInFile;
				sErrMsg += "\n";
				return false;
			}
			parser = new yard::SimpleTextParser(data, data + datasize);	   
			if (!parser->Parse<heron_grammar::HeronFile>()) {
				sErrMsg += "Failed to parse file: ";
				sErrMsg += sInFile;
				sErrMsg += "\n";
				sErrMsg += parser->OutputLocation();
				printf("%s\n", sErrMsg.c_str());
				return false;
			}
			try {	
				root = new HeronFile(NULL, parser->GetAstRoot());
			}
			catch (AstException e) {
				sErrMsg += "Error occured during AST construction: ";
				sErrMsg += e.msg;
				sErrMsg += "\n";
				return false;
			}
			sErrMsg = "success\n";
			return true;
		}

		ConstString GetErrorMsg() {
			return sErrMsg;
		}

		HeronFile* GetRoot() {
			return root;
		}
	};
}

#endif
